import pathlib

import matplotlib.pyplot
import torch

from checkpoint import Checkpoint
from checkpoint_manager import CheckpointManager
from scaler import Scaler
from simple_pendulum_dataset import SimplePendulumDataset
from spinn import Spinn


def main():
    problems = [
        (1, "not_damped", Spinn()),
        (2, "under_damping", Spinn()),
        (3, "over_damping", Spinn(30, 1))
    ]

    for subplot_index, title, model in problems:
        model_path = pathlib.Path(f"./checkpoints/{title}")
        data_path = pathlib.Path(f"./data/{title}.pt")

        dataset = SimplePendulumDataset.load(data_path)
        t_min = float(dataset.t_min())
        t_max = float(dataset.t_max())
        model.scalar = Scaler(t_min, t_max)

        checkpoints = CheckpointManager(model_path).existed()
        checkpoints = sorted(checkpoints.items(), key=lambda x: x[0])
        checkpoints_step = len(checkpoints) // 100
        if checkpoints_step > 0:
            checkpoints = checkpoints[checkpoints_step - 1::checkpoints_step]

        matplotlib.pyplot.subplot(2, len(problems), subplot_index)
        line_g = matplotlib.pyplot.axhline(dataset.problem().g(), c="tab:green")

        matplotlib.pyplot.subplot(2, len(problems), subplot_index + len(problems))
        line_gamma = matplotlib.pyplot.axhline(dataset.problem().gamma(), c="tab:green")

        scatter_g = None
        scatter_gamma = None
        for epoch, checkpoint_path in checkpoints:
            checkpoint = Checkpoint.load(checkpoint_path)
            model.load_state_dict(checkpoint.model_state())

            matplotlib.pyplot.subplot(2, len(problems), subplot_index)
            scatter_g = matplotlib.pyplot.scatter(
                epoch, torch.detach(model.g),
                c="tab:blue", s=4)

            matplotlib.pyplot.subplot(2, len(problems), subplot_index + len(problems))
            scatter_gamma = matplotlib.pyplot.scatter(
                epoch, torch.detach(model.gamma),
                c="tab:orange", s=4)

        matplotlib.pyplot.subplot(2, len(problems), subplot_index)
        matplotlib.pyplot.title(title)
        matplotlib.pyplot.ylabel(r"$g$")
        matplotlib.pyplot.legend(
            [line_g, scatter_g],
            [f"Expected = {dataset.problem().g():.2f}", f"Estimated = {float(model.g):.2f}"])

        matplotlib.pyplot.subplot(2, len(problems), subplot_index + len(problems))
        matplotlib.pyplot.xlabel("Epoch")
        matplotlib.pyplot.ylabel(r"$\gamma$")
        matplotlib.pyplot.legend(
            [line_gamma, scatter_gamma],
            [f"Expected = {dataset.problem().gamma():.2f}", f"Estimated = {float(model.gamma):.2f}"])

    matplotlib.pyplot.subplots_adjust(wspace=0.25)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
